#!/usr/bin/env python3
"""
SAM检查点下载脚本

自动下载SAM模型检查点文件
"""

import os
import sys
import urllib.request
from pathlib import Path
from tqdm import tqdm


def download_file(url: str, filepath: Path, desc: str = "Downloading"):
    """下载文件并显示进度"""
    filepath.parent.mkdir(parents=True, exist_ok=True)
    
    if filepath.exists():
        print(f"文件已存在: {filepath}")
        return True
    
    try:
        with urllib.request.urlopen(url) as response:
            total_size = int(response.headers.get('Content-Length', 0))
            
            with open(filepath, 'wb') as f:
                with tqdm(total=total_size, unit='B', unit_scale=True, desc=desc) as pbar:
                    while True:
                        chunk = response.read(8192)
                        if not chunk:
                            break
                        f.write(chunk)
                        pbar.update(len(chunk))
        
        print(f"下载完成: {filepath}")
        return True
    except Exception as e:
        print(f"下载失败: {e}")
        return False


def download_sam_checkpoints():
    """下载SAM检查点文件"""
    project_root = Path(__file__).parent.parent
    checkpoint_dir = project_root / 'checkpoints' / 'sam'
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    
    # SAM检查点URL（官方GitHub发布）
    checkpoints = {
        'sam_vit_h': {
            'url': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth',
            'path': checkpoint_dir / 'sam_vit_h_4b8939.pth',
            'size': '2.4GB'
        },
        'sam_vit_l': {
            'url': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth',
            'path': checkpoint_dir / 'sam_vit_l_0b3195.pth',
            'size': '1.2GB'
        },
        'sam_vit_b': {
            'url': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth',
            'path': checkpoint_dir / 'sam_vit_b_01ec64.pth',
            'size': '375MB'
        }
    }
    
    print("=" * 60)
    print("SAM检查点下载工具")
    print("=" * 60)
    print("\n可用模型:")
    for i, (name, info) in enumerate(checkpoints.items(), 1):
        exists = "✓" if info['path'].exists() else " "
        print(f"{i}. {name} ({info['size']}) {exists}")
    
    print("\n请选择要下载的模型（输入数字，多个用逗号分隔，或输入'all'下载全部）:")
    choice = input("> ").strip().lower()
    
    if choice == 'all':
        models_to_download = list(checkpoints.keys())
    else:
        try:
            indices = [int(x.strip()) - 1 for x in choice.split(',')]
            models_to_download = [list(checkpoints.keys())[i] for i in indices if 0 <= i < len(checkpoints)]
        except:
            print("无效输入，退出")
            return
    
    if not models_to_download:
        print("未选择任何模型，退出")
        return
    
    print(f"\n开始下载 {len(models_to_download)} 个模型...")
    
    for model_name in models_to_download:
        info = checkpoints[model_name]
        if info['path'].exists():
            print(f"\n跳过 {model_name}（已存在）")
            continue
        
        print(f"\n下载 {model_name} ({info['size']})...")
        success = download_file(info['url'], info['path'], desc=f"下载 {model_name}")
        
        if not success:
            print(f"警告: {model_name} 下载失败")
    
    print("\n" + "=" * 60)
    print("下载完成！")
    print(f"检查点保存在: {checkpoint_dir}")
    print("=" * 60)


if __name__ == "__main__":
    download_sam_checkpoints()

